{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook shows how to use the simple model serving functionality found in `tfe.serving`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-11-01 10:11:33.040484: I tensorflow/core/util/util.cc:169] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2022-11-01 10:11:33.044530: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory\n",
      "2022-11-01 10:11:33.044542: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.\n",
      "2022-11-01 10:11:34.342585: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory\n",
      "2022-11-01 10:11:34.342610: W tensorflow/stream_executor/cuda/cuda_driver.cc:269] failed call to cuInit: UNKNOWN ERROR (303)\n",
      "2022-11-01 10:11:34.342623: I tensorflow/stream_executor/cuda/cuda_diagnostics.cc:156] kernel driver does not appear to be running on this host (mpc-empty011122134151.ea134): /proc/driver/nvidia/version does not exist\n",
      "2022-11-01 10:11:34.342859: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 AVX512F AVX512_VNNI FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "import tf_encrypted as tfe\n",
    "\n",
    "from collections import OrderedDict"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Protocol\n",
    "\n",
    "We first configure the protocol we will be using, as well as the servers on which we want to run it.\n",
    "\n",
    "Note that the configuration is saved to file as we will be needing it in the client as well."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-11-01 10:11:34.412635: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272] Initialize GrpcChannelCache for job tfe -> {0 -> localhost:4000, 1 -> localhost:4001, 2 -> localhost:4002}\n",
      "2022-11-01 10:11:34.412659: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272] Initialize GrpcChannelCache for job localhost -> {0 -> localhost:44660}\n",
      "2022-11-01 10:11:34.420067: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272] Initialize GrpcChannelCache for job tfe -> {0 -> localhost:4000, 1 -> localhost:4001, 2 -> localhost:4002}\n",
      "2022-11-01 10:11:34.420085: I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:272] Initialize GrpcChannelCache for job localhost -> {0 -> localhost:44660}\n",
      "2022-11-01 10:11:34.420396: I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:438] Started server with target: grpc://localhost:44660\n"
     ]
    }
   ],
   "source": [
    "players = OrderedDict([\n",
    "    ('server0', 'localhost:4000'),\n",
    "    ('server1', 'localhost:4001'),\n",
    "    ('server2', 'localhost:4002'),\n",
    "])\n",
    "\n",
    "config = tfe.RemoteConfig(players)\n",
    "config.connect_servers()\n",
    "tfe.set_config(config)\n",
    "tfe.set_protocol(tfe.protocol.Pond())\n",
    "config.save('./tfe.config')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Launching servers\n",
    "\n",
    "Before actually serving the computation below we need to launch TensorFlow servers in new processes. Run the following in three different terminals. You may have to allow Python to accept incoming connections."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "python -m tf_encrypted.player --config ./tfe.config server0\n",
      "python -m tf_encrypted.player --config ./tfe.config server1\n",
      "python -m tf_encrypted.player --config ./tfe.config server2\n"
     ]
    }
   ],
   "source": [
    "for player_name in players.keys():\n",
    "    print(\"python -m tf_encrypted.player --config ./tfe.config {}\".format(player_name))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computation\n",
    "\n",
    "We then define the computation we want to run. These will happen on private tensors on the servers defined above.\n",
    "\n",
    "Note that the only single-tensor computations are currently supported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_shape = (5, 5)\n",
    "output_shape = (5, 5)\n",
    "\n",
    "def computation(x):\n",
    "    return x * x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can then set up a new `tfe.serving.QueueServer` to serve this computation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "server = tfe.serving.QueueServer(\n",
    "    input_shape=input_shape,\n",
    "    output_shape=output_shape,\n",
    "    computation_fn=computation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Serving\n",
    "\n",
    "With all of the above in place we can finally connect to our servers, push our graph to them, and start serving computations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Next\n"
     ]
    }
   ],
   "source": [
    "def step_fn():\n",
    "    print(\"Next\")\n",
    "\n",
    "server.run(num_steps=1, step_fn=step_fn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point we switch to the client notebook to run computations."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.13 ('tfe2')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "2b3cf1a821f25a2bafd97d4b4f3a77d0e1f42d8d78c0e512d3a72b703c70416e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
